#include "native/sun/misc/Unsafe.h"
#include "rtda/operand_stack.h"
#include "rtda/local_vars.h"
#include <algorithm>

namespace jvm {
    namespace misc {

        std::map<int64_t, std::vector<uint8_t>> Unsafe::g_allocatedMemory;
        int64_t Unsafe::g_nextAddress = 64;

        void Unsafe::registerNatives() {
            native::Registry::registerMethod("sun/misc/Unsafe", "arrayBaseOffset",
                "(Ljava/lang/Class;)I", arrayBaseOffset);
            native::Registry::registerMethod("sun/misc/Unsafe", "arrayIndexScale",
                "(Ljava/lang/Class;)I", arrayIndexScale);
            native::Registry::registerMethod("sun/misc/Unsafe", "addressSize",
                "()I", addressSize);
            native::Registry::registerMethod("sun/misc/Unsafe", "objectFieldOffset",
                "(Ljava/lang/reflect/Field;)J", objectFieldOffset);
            native::Registry::registerMethod("sun/misc/Unsafe", "compareAndSwapObject",
                "(Ljava/lang/Object;JLjava/lang/Object;Ljava/lang/Object;)Z", compareAndSwapObject);
            native::Registry::registerMethod("sun/misc/Unsafe", "getIntVolatile",
                "(Ljava/lang/Object;J)I", getInt);
            native::Registry::registerMethod("sun/misc/Unsafe", "compareAndSwapInt",
                "(Ljava/lang/Object;JII)Z", compareAndSwapInt);
            native::Registry::registerMethod("sun/misc/Unsafe", "getObjectVolatile",
                "(Ljava/lang/Object;J)Ljava/lang/Object;", getObject);
            native::Registry::registerMethod("sun/misc/Unsafe", "compareAndSwapLong",
                "(Ljava/lang/Object;JJJ)Z", compareAndSwapLong);
        }

        void Unsafe::arrayBaseOffset(std::shared_ptr<rtda::Frame> frame) {
            frame->getOperandStack()->pushInt(0); // TODO
        }

        void Unsafe::arrayIndexScale(std::shared_ptr<rtda::Frame> frame) {
            frame->getOperandStack()->pushInt(1); // TODO
        }

        void Unsafe::addressSize(std::shared_ptr<rtda::Frame> frame) {
            frame->getOperandStack()->pushInt(8); // TODO
        }

        void Unsafe::objectFieldOffset(std::shared_ptr<rtda::Frame> frame) {
            rtda::LocalVars& vars = frame->getLocalVars();
            std::shared_ptr<heap::Object> jField = vars.getRef(1);

            int32_t offset = jField->getIntVar("slot", "I");
            frame->getOperandStack()->pushLong(offset);
        }

        void Unsafe::compareAndSwapObject(std::shared_ptr<rtda::Frame> frame) {
            rtda::LocalVars& vars = frame->getLocalVars();
            std::shared_ptr<heap::Object> obj = vars.getRef(1);
            void* fields = obj->data();
            int64_t offset = vars.getLong(2);
            std::shared_ptr<heap::Object> expected = vars.getRef(4);
            std::shared_ptr<heap::Object> newVal = vars.getRef(5);

            if (auto slots = dynamic_cast<rtda::Slots*>(fields)) {
                bool swapped = casObject(obj, *slots, offset, expected, newVal);
                frame->getOperandStack().pushBoolean(swapped);
            }
            else if (auto objs = dynamic_cast<std::vector<std::shared_ptr<heap::Object>>*>(fields)) {
                bool swapped = casArray(*objs, offset, expected, newVal);
                frame->getOperandStack().pushBoolean(swapped);
            }
            else {
                throw std::runtime_error("TODO: compareAndSwapObject!");
            }
        }

        bool Unsafe::casObject(std::shared_ptr<heap::Object> obj, rtda::Slots& fields, int64_t offset,
            std::shared_ptr<heap::Object> expected, std::shared_ptr<heap::Object> newVal) {
            std::shared_ptr<heap::Object> current = fields.getRef(static_cast<uint32_t>(offset));
            if (current == expected) {
                fields.setRef(static_cast<uint32_t>(offset), newVal);
                return true;
            }
            return false;
        }

        bool Unsafe::casArray(std::vector<std::shared_ptr<heap::Object>>& objs, int64_t offset,
            std::shared_ptr<heap::Object> expected, std::shared_ptr<heap::Object> newVal) {
            std::shared_ptr<heap::Object> current = objs[offset];
            if (current == expected) {
                objs[offset] = newVal;
                return true;
            }
            return false;
        }

        void Unsafe::getInt(std::shared_ptr<rtda::Frame> frame) {
            rtda::LocalVars& vars = frame->getLocalVars();
            void* fields = vars.getRef(1)->data();
            int64_t offset = vars.getLong(2);

            if (auto slots = dynamic_cast<rtda::Slots*>(fields)) {
                frame->getOperandStack().pushInt(slots->getInt(static_cast<uint32_t>(offset)));
            }
            else if (auto ints = dynamic_cast<std::vector<int32_t>*>(fields)) {
                frame->getOperandStack().pushInt((*ints)[offset]);
            }
            else {
                throw std::runtime_error("getInt!");
            }
        }

        // Other Unsafe method implementations...

        int64_t Unsafe::allocateMemory(int64_t bytes) {
            std::vector<uint8_t> mem(bytes);
            int64_t address = g_nextAddress;
            g_allocatedMemory[address] = std::move(mem);
            g_nextAddress += bytes;
            return address;
        }

        int64_t Unsafe::reallocateMemory(int64_t address, int64_t bytes) {
            if (bytes == 0) return 0;
            if (address == 0) return allocateMemory(bytes);

            auto& mem = memoryAt(address);
            if (mem.size() >= static_cast<size_t>(bytes)) {
                return address;
            }

            g_allocatedMemory.erase(address);
            int64_t newAddress = allocateMemory(bytes);
            auto& newMem = memoryAt(newAddress);
            std::copy(mem.begin(), mem.end(), newMem.begin());
            return newAddress;
        }

        void Unsafe::freeMemory(int64_t address) {
            if (g_allocatedMemory.erase(address) == 0) {
                throw std::runtime_error("Memory was not allocated!");
            }
        }

        std::vector<uint8_t>& Unsafe::memoryAt(int64_t address) {
            for (auto& [startAddress, mem] : g_allocatedMemory) {
                int64_t endAddress = startAddress + static_cast<int64_t>(mem.size());
                if (address >= startAddress && address < endAddress) {
                    int64_t offset = address - startAddress;
                    return mem[offset];
                }
            }
            throw std::runtime_error("Invalid address!");
        }

    } // namespace misc
} // namespace jvm